import collections

# 状态机
# class Solution(object):
#     def candy(self, ratings):
#         """
#         :type ratings: List[int]
#         :rtype: int
#         """
#         length = len(ratings)
#         left_index = 0
#         right_index = 1
#         state = 0
#         Q = collections.deque()
#         award = [0 for _ in range(length)]
#         while right_index < length:
#             if ratings[right_index] > ratings[right_index - 1]:
#                 if state == 1:
#                     right_index += 1
#                 else:
#                     if state == 0:
#                         left_index = right_index - 1
#                         right_index += 1
#                         state = 1
#                     else:
#                         Q.append([left_index, right_index - 1])
#                         left_index = right_index
#                         right_index += 1
#                         state = 1
#             else:
#                 if ratings[right_index] < ratings[right_index - 1]:
#                     if state == 0:
#                         left_index = right_index - 1
#                         right_index = right_index + 1
#                         state = 2
#                     else:
#                         if state == 1:
#                             Q.append([left_index, right_index - 1])
#                             left_index = right_index
#                             right_index += 1
#                             state = 2
#                         else:
#                             right_index += 1
#                 else:
#                     if state == 0:
#                         left_index = right_index
#                         right_index += 1
#                     else:
#                         if state == 1:
#                             Q.append([left_index, right_index - 1])
#                             left_index = right_index
#                             right_index += 1
#                             state = 0
#                         else:
#                             Q.append([left_index, right_index - 1])
#                             left_index = right_index
#                             right_index += 1
#                             state = 0
#         if state == 1:
#             Q.append([left_index, right_index - 1])
#         if state == 2:
#             Q.append([left_index, right_index - 1])
#         state = 0
#         last_index = -1
#         for item in Q:
#             if item[0] > last_index + 1:
#                 for index in range(last_index + 1, item[0]):
#                     award[index] = 1
#                 state = 0
#                 last_index = item[0] - 1
#             if last_index >= 0 and ratings[item[0]] == ratings[last_index]:
#                 state = 0
#             if state == 0:
#                 if ratings[item[1]] > ratings[item[0]]:
#                     for num, index in enumerate(range(item[0], item[1] + 1)):
#                         award[index] = num + 1
#                     state = 1
#                 else:
#                     for num, index in enumerate(range(item[1], item[0] - 1, -1)):
#                         award[index] = num + 1
#                     state = 2
#                 last_index = item[1]
#             else:
#                 if state == 1:
#                     if item[1] == item[0]:
#                         award[item[0]] = 1
#                         last_index = item[0]
#                     else:
#                         for num, index in enumerate(range(item[1], item[0] - 1, -1)):
#                             award[index] = num + 1
#                         if award[item[0]] >= award[last_index]:
#                             award[last_index] = award[item[0]] + 1
#                         last_index = item[1]
#                     state = 2
#                 else:
#                     if item[1] == item[0]:
#                         award[item[0]] = 2
#                         last_index = item[0]
#                     else:
#                         for num, index in enumerate(range(item[0], item[1] + 1)):
#                             award[index] = num + 2
#                         last_index = item[1]
#                     state = 1
#         if last_index + 1 < length:
#             for i in range(last_index + 1, length):
#                 award[i] = 1
#         # for i in range(len(award)):
#         #     print(i, sum(award[: i]))
#         return sum(award)
#
# data = Solution()
# ratings = [169, 311, 517, 4121, 3946, 3946]
# print(ratings)
# print(data.candy(ratings))

# 单调栈
class Solution(object):
    def candy(self, ratings):
        """
        :type ratings: List[int]
        :rtype: int
        """
        last_value = 0
        stack = collections.deque([0])
        index = 1
        total_num = 0
        last_max_value = 1
        while index < len(ratings):
            if ratings[index] > ratings[stack[-1]]:
                if len(stack) > 1:
                    last_value = 0
                    while stack:
                        last_value += 1
                        total_num += last_value
                        stack.pop()
                    total_num += max(last_max_value - last_value, 0)
                    last_max_value = 1
                    last_value = 1
                else:
                    last_value += 1
                    stack.pop()
                    total_num += last_value
            else:
                if ratings[index] == ratings[stack[-1]]:
                    if len(stack) > 1:
                        last_value = 0
                        while stack:
                            last_value += 1
                            total_num += last_value
                            stack.pop()
                        total_num += max(last_max_value - last_value, 0)
                        last_max_value = 1
                        last_value = 0
                    else:
                        stack.pop()
                        total_num += (last_value + 1)
                        last_value = 0
                else:
                    if len(stack) == 1:
                        last_max_value = last_value + 1
            stack.append(index)
            index += 1
        if stack:
            if len(stack) > 1:
                last_value = 0
                while stack:
                    last_value += 1
                    total_num += last_value
                    stack.pop()
                total_num += max(last_max_value - last_value, 0)
            else:
                total_num += (last_value + 1)
                stack.pop()
        return total_num

data = Solution()
ratings =[1, 2, 3, 4, 1, 1]
print(data.candy(ratings))
